
# 提取文本和图片
# 支持doc、docx、pdf
# 输出json文件，包含文本和图片路径
# 图片保存到images目录下，以pdf名称命名文件夹
# 文本保存到pdf或word目录下，以pdf或word名称命名文件
# 图片和文本在json中以列表形式存在
# 输出结果样式如下
# {
#     "extracted_text": "文本内容 \n[图片路径1]\n文本内容\n[图片路径2]\n文本内容",
#     "source_file": "文件名",
#     "images": ["图片路径1", "图片路径2", ...]
# }

import json
from pathlib import Path
import docx
import PyPDF2
import textract
import os
import sys
import fitz  # PyMuPDF
from PyQt6.QtWidgets import QApplication, QFileDialog


def extract_from_doc(file_path: str) -> str:
    """从旧版 Word (.doc) 文档中提取文本"""
    try:
        text = textract.process(file_path).decode('utf-8')
        return text.strip()
    except Exception as e:
        raise Exception(f"处理 .doc 文件时出错: {str(e)}")


def extract_from_docx(file_path: str) -> str:
    """从 Word (.docx) 文档中提取文本"""
    doc = docx.Document(file_path)
    return '\n'.join([paragraph.text for paragraph in doc.paragraphs])


class TextExtractor:
    def __init__(self):
        """初始化提取器，创建输出目录"""
        # 设置基础输出目录
        base_dir = Path("/Volumes/HIKVISION/0、学习/聊天json/【男生版】让女生追着你聊天的话术技巧/word_pdf_extract/output")

        # 设置各类文件的输出目录
        self.pdf_output_dir = base_dir / "pdf"
        self.word_output_dir = base_dir / "word"
        self.images_output_dir = base_dir / "images"  # 添加图片输出目录

        # 确保所有输出目录存在
        for dir_path in [self.pdf_output_dir, self.word_output_dir, self.images_output_dir]:
            dir_path.mkdir(parents=True, exist_ok=True)

    def extract_images_from_pdf(self, file_path: str) -> list:
        """从PDF文档中提取图片"""
        image_list = []
        pdf_file = fitz.open(file_path)
        pdf_name = Path(file_path).stem
        images_dir = self.images_output_dir / pdf_name
        images_dir.mkdir(exist_ok=True)

        try:
            for page_index in range(len(pdf_file)):
                page = pdf_file[page_index]
                image_list.extend(page.get_images())

            if image_list:
                # 提取图片
                for image_index, img in enumerate(image_list):
                    xref = img[0]
                    base_image = pdf_file.extract_image(xref)
                    image_bytes = base_image["image"]
                    image_ext = base_image["ext"]
                    image_path = images_dir / f"image_{image_index + 1}.{image_ext}"

                    with open(image_path, "wb") as image_file:
                        image_file.write(image_bytes)
                        print(f"保存图片: {image_path}")

            return [str(images_dir / f"image_{i + 1}.{pdf_file.extract_image(img[0])['ext']}")
                    for i, img in enumerate(image_list)]

        finally:
            pdf_file.close()

    def extract_from_pdf(self, file_path: str) -> tuple:
        """从PDF文档中提取文本和图片，并在文本中标注图片位置"""
        pdf_file = fitz.open(file_path)
        text_parts = []
        image_paths = []

        try:
            for page_index in range(len(pdf_file)):
                page = pdf_file[page_index]

                # 获取页面文本和图片
                text_blocks = page.get_text("blocks")  # 获取文本块
                images = page.get_images()

                # 如果页面有图片，需要处理图文混排
                if images:
                    # 获取每个图片的位置信息
                    image_info = []
                    for img_index, img in enumerate(images):
                        xref = img[0]
                        base_image = pdf_file.extract_image(xref)
                        image_ext = base_image["ext"]
                        image_filename = f"image_{page_index + 1}_{img_index + 1}.{image_ext}"

                        # 保存图片
                        pdf_name = Path(file_path).stem
                        images_dir = self.images_output_dir / pdf_name
                        images_dir.mkdir(exist_ok=True)
                        image_path = images_dir / image_filename

                        with open(image_path, "wb") as image_file:
                            image_file.write(base_image["image"])
                            print(f"保存图片: {image_path}")
                            image_paths.append(str(image_path))

                        # 记录图片位置和文件名
                        image_info.append({
                            'filename': image_filename,
                            'y0': img[3],  # 图片的垂直位置
                            'path': str(image_path)
                        })

                    # 按垂直位置排序文本块和图片
                    text_blocks = [b for b in text_blocks if b[6] == 0]  # 只保留文本块
                    text_blocks.sort(key=lambda x: x[1])  # 按y坐标排序

                    # 处理文本块，在适当位置插入图片标记
                    for i in range(len(text_blocks) - 1):
                        current_block = text_blocks[i]
                        next_block = text_blocks[i + 1]

                        # 添加当前文本块
                        text_parts.append(current_block[4])

                        # 检查是否有图片在当前文本块和下一个文本块之间
                        current_y_max = current_block[3]  # 当前块的底部y坐标
                        next_y_min = next_block[1]  # 下一块的顶部y坐标

                        # 查找位于两个文本块之间的图片
                        for img in image_info:
                            if current_y_max <= img['y0'] <= next_y_min:
                                text_parts.append(f"\n[{img['filename']}]\n")

                    # 添加最后一个文本块
                    if text_blocks:
                        text_parts.append(text_blocks[-1][4])
                else:
                    # 如果页面没有图片，直接添加文本
                    text = page.get_text()
                    if text.strip():
                        text_parts.append(text.strip())

        finally:
            pdf_file.close()

        # 合并所有文本部分
        final_text = "\n".join(text_parts)
        return text_parts, image_paths  # 不再合并文本

    def save_to_json(self, text: str, original_file: Path, is_pdf: bool, image_paths: list = None) -> None:
        """将提取的文本和图片路径保存为JSON文件"""
        global text_lines
        output_dir = self.pdf_output_dir if is_pdf else self.word_output_dir
        output_path = output_dir / f"{original_file.stem}.json"

        if isinstance(text, tuple):
            text = text[0]

        # 处理文本，将其分割成行
        if isinstance(text, str):
            text_lines = [line.strip() for line in text.split('\n') if line.strip()]
        elif isinstance(text, tuple):
            text_lines = [line.strip() for line in text[0].split('\n') if line.strip()]
        elif isinstance(text, list):
            # 如果是列表，将每个元素按换行符分割，然后展平列表
            text_lines = []
            for item in text:
                if isinstance(item, str):
                    lines = [line.strip() for line in item.split('\n') if line.strip()]
                    text_lines.extend(lines)

        data = {
            'extracted_text': {
                'texts': text_lines,  # 确保文本是列表形式
                'pictures': image_paths if image_paths else []
            },
            'source_file': original_file.name,
        }

        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
            print(f"文件已保存到: {output_path}")

    def process_file(self, file_path: str) -> None:
        """处理单个文件并保存其内容"""
        path = Path(file_path).resolve()  # 获取绝对路径

        try:
            if not path.exists():
                raise FileNotFoundError(f"文件不存在: {path}")

            print(f"正在处理文件: {path.name}")

            if path.suffix.lower() == '.docx':
                text = extract_from_docx(str(path))
                self.save_to_json(text, path, False)
            elif path.suffix.lower() == '.doc':
                text = extract_from_doc(str(path))
                self.save_to_json(text, path, False)
            elif path.suffix.lower() == '.pdf':
                text = self.extract_from_pdf(str(path))
                self.save_to_json(text, path, True)
            else:
                raise ValueError(f"不支持的文件类型: {path.suffix}")

            print(f"成功处理文件 {path.name}")

        except Exception as e:
            print(f"处理文件 {path.name} 时出错: {str(e)}")

def process_directory(directory: str):
    """处理指定目录下的所有文档文件"""
    extractor = TextExtractor()
    dir_path = Path(directory)

    if not dir_path.exists():
        print(f"目录不存在: {directory}")
        return

    for file_path in dir_path.glob('*.*'):
        if file_path.suffix.lower() in ['.doc', '.docx', '.pdf']:
            extractor.process_file(str(file_path))


def select_input_files():
    """打开文件选择对话框，支持多选文件"""
    app = QApplication(sys.argv)

    # 设置默认打开路径
    default_path = "/Volumes/HIKVISION/0、学习/聊天/【男生版】让女生追着你聊天的话术技巧"

    file_dialog = QFileDialog()
    file_dialog.setFileMode(QFileDialog.FileMode.ExistingFiles)
    file_dialog.setNameFilter("文档文件 (*.doc *.docx *.pdf)")
    file_dialog.setDirectory(default_path)  # 设置默认目录

    if file_dialog.exec():
        return file_dialog.selectedFiles()
    return []


# # 使用代码路径
# if __name__ == "__main__":
#     extractor = TextExtractor()
#     extractor.process_file("/Volumes/HIKVISION/0、学习/聊天/【男生版】让女生追着你聊天的话术技巧/7、与m 聊天的128个步步深入话题.doc")
#     # extractor.process_file("example.docx")
#     # extractor.process_file("example.pdf")

# 使用文件选择对话框选择文件
if __name__ == "__main__":
    # 打开文件选择对话框
    selected_files = select_input_files()

    if selected_files:  # 如果用户选择了文件
        extractor = TextExtractor()
        total_files = len(selected_files)

        print(f"共选择了 {total_files} 个文件")

        for index, file_path in enumerate(selected_files, 1):
            print(f"\n处理第 {index}/{total_files} 个文件:")
            try:
                extractor.process_file(file_path)
            except Exception as e:
                print(f"处理文件失败: {str(e)}")

        print("\n所有文件处理完成！")
    else:
        print("未选择任何文件")